import pygame
from plane import Plane
from obstacle import Obstacle
from target import Target
import numpy as np
from threading import Thread

# 环境
class MyEnv:
    def __init__(self, screen_size=(800,800), plane_name=1, obstacle_num=3, target_num=1) -> None:
        self.screen_size = screen_size

        self.plane_num = plane_name
        self.obstacle_num = obstacle_num
        self.target_num = target_num

        # 创建物体
        self.planes = [Plane() for _ in range(self.plane_num)]
        self.obstacles = [Obstacle() for _ in range(self.obstacle_num)]
        self.targets = [Target() for _ in range(self.target_num)]

        pygame.init()
        self.screen = pygame.display.set_mode(self.screen_size)
        pygame.display.set_caption("my env")         
 
    # 重置环境
    def reset(self):
        # 随机出生位置
        for plane in self.planes:
            plane.rect.center = (np.random.randint(0, self.screen_size[0]), np.random.randint(0, self.screen_size[1]))
        for obstacle in self.obstacles:
            obstacle.rect.center = (np.random.randint(0, self.screen_size[0]), np.random.randint(0, self.screen_size[1]))
        for target in self.targets:
            target.rect.center = (np.random.randint(0, self.screen_size[0]), np.random.randint(0, self.screen_size[1]))

        # 计算plane距离target的距离
        pos_plane = self.planes[0].rect.center
        pos_target = self.targets[0].rect.center
        self.last_dis = d(pos_plane, pos_target)

        return self.get_obs()

    # 环境行动一步
    def step(self, actions):
        # 飞机移动
        for i in range(self.plane_num):
            self.planes[i].step(actions[i], self.screen_size)

        # 重新绘制环境
        self.screen.fill((255,255,255))
        self.draw()
        pygame.display.flip()

        # 事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
        
        return self.get_obs(), self.get_reward(), self.is_done()
    
    # 绘制所有物体
    def draw(self):
        for plane in self.planes:
            plane.draw(self.screen)
        for obstacle in self.obstacles:
            obstacle.draw(self.screen)
        for target in self.targets:
            target.draw(self.screen)
    
    # 获取当前状态
    def get_obs(self):
        # 当前状态
        obs = []
        for thing in self.planes + self.obstacles + self.targets:
            obs.append(thing.rect.centerx - self.planes[0].rect.centerx)
            obs.append(thing.rect.centery - self.planes[0].rect.centery)
        return obs

    # 获取奖励
    def get_reward(self):
        to_obstacle, to_target = self.check_collision()

        # 离目标距离惩罚
        pos_plane = self.planes[0].rect.center
        pos_target = self.targets[0].rect.center
        now_dis = d(pos_plane, pos_target)
        reward_target = -10 * now_dis
        
        # 靠近障碍物惩罚
        reward_obs = 0
        for obs in self.obstacles:
            pos_obs = obs.rect.center
            if d(pos_plane, pos_obs) <= 30:
                reward_obs += -1 / 2 * (1 / d(pos_plane, pos_obs) - 1 / 30) ** 2
        
        # 远离目标惩罚
        reward_away = 100 * (self.last_dis - now_dis)
        self.last_dis = now_dis
        
        # 撞上障碍物惩罚
        reward_arrive = 10000 if to_target else 0
        
        # 达到目标奖励
        reward_collision = -10000 if to_obstacle else 0
        
        return reward_target + reward_obs + reward_away + reward_arrive + reward_collision

    # 是否结束
    def is_done(self):
        to_obstacle, to_target = self.check_collision()
        return to_obstacle or to_target

    # 碰撞检测
    def check_collision(self):
        to_obstacle = False  # 是否撞上障碍物
        to_target = False  # 是否到达目标
        for p in self.planes:
            for o in self.obstacles:
                if pygame.sprite.collide_mask(p, o):
                    to_obstacle = True
                    break
            if to_obstacle:
                break;
        for p in self.planes:
            for t in self.targets:
                if pygame.sprite.collide_mask(p, t):
                    to_target = True
                    break
            if to_target:
                break
        
        return to_obstacle, to_target
    
    # 随机一个动作
    def sample_action(self):
        return np.random.random(1) * 2 * np.pi - np.pi

# 计算两个位置曼哈顿距离
def d(p1, p2):
        return np.sum(np.abs(np.array(p1) - np.array(p2)))